from typing import TypedDict, Annotated, Optional, List, Any
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import AnyMessage
from langgraph.graph import add_messages



class Nl2SqlState(TypedDict):
    nl2sql_messages: Annotated[list[AnyMessage], add_messages]
    db: Optional[SQLDatabase]
    analysis_model: Optional[str]
    main_table: Optional[str]
    table_schema: Optional[str]
    configuration_tables: Optional[List[str]]
    configuration_tables_schema: Optional[List[str]]
    sql: Optional[str]
    description: Optional[List[str]]





class QueryOutput(TypedDict):
    """生成的SQL查询语句"""

    sql: Annotated[str, ..., "Syntactically valid SQL query."]


class TableOutput(TypedDict):
    """生成的SQL查询语句"""

    table_name: Annotated[str, ..., "main table name"]


